package my.shadowsocks.common.encryption;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.security.InvalidAlgorithmParameterException;
import java.security.SecureRandom;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

import javax.crypto.SecretKey;
import javax.crypto.spec.SecretKeySpec;

import org.bouncycastle.crypto.StreamBlockCipher;
import org.bouncycastle.crypto.params.KeyParameter;
import org.bouncycastle.crypto.params.ParametersWithIV;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * @author fengbo
 */
public abstract class AbstractCrypt implements Crypt {

	protected final String name;
	private final SecretKey key;
	private final int ivLength;
	private boolean encryptIVSet;
	private boolean decryptIVSet;
	private final Lock encLock = new ReentrantLock();
	private final Lock decLock = new ReentrantLock();
	protected StreamBlockCipher encCipher;
	protected StreamBlockCipher decCipher;

	private static final Logger logger = LoggerFactory.getLogger(AbstractCrypt.class);

	public AbstractCrypt(String name, String password) {
		this.name = name.toLowerCase();
		this.ivLength = getIVLength();
		ShadowSocksKey ssKey = new ShadowSocksKey(password, getKeyLength());
		this.key = new SecretKeySpec(ssKey.getEncoded(), "AES");
	}

	private void setIV(byte[] iv, boolean isEncrypt) {
		if (ivLength == 0) {
			return;
		}

		if (isEncrypt) {
			byte[] encryptIV = new byte[ivLength];
			System.arraycopy(iv, 0, encryptIV, 0, ivLength);
			try {
				encCipher = getCipher(true);
				encCipher.init(true, new ParametersWithIV(new KeyParameter(key.getEncoded()), encryptIV));
			} catch (InvalidAlgorithmParameterException e) {
				logger.info(e.toString());
			}
		} else {
			byte[] decryptIV = new byte[ivLength];
			System.arraycopy(iv, 0, decryptIV, 0, ivLength);
			try {
				decCipher = getCipher(false);
				decCipher.init(false, new ParametersWithIV(new KeyParameter(key.getEncoded()), decryptIV));
			} catch (InvalidAlgorithmParameterException e) {
				logger.info(e.getMessage(), e);
			}
		}
	}

	@Override
	public void encrypt(byte[] data, ByteArrayOutputStream stream) {
		synchronized (encLock) {
			stream.reset();
			if (!encryptIVSet) {
				encryptIVSet = true;
				byte[] iv = randomBytes(ivLength);
				setIV(iv, true);
				try {
					stream.write(iv);
				} catch (IOException e) {
					logger.info(e.toString());
				}
			}
			encrypt0(data, stream);
		}
	}

	@Override
	public void encrypt(byte[] data, int length, ByteArrayOutputStream stream) {
		byte[] d = new byte[length];
		System.arraycopy(data, 0, d, 0, length);
		encrypt(d, stream);
	}

	@Override
	public void decrypt(byte[] data, ByteArrayOutputStream stream) {
		byte[] temp;
		synchronized (decLock) {
			stream.reset();
			if (!decryptIVSet) {
				decryptIVSet = true;
				setIV(data, false);
				temp = new byte[data.length - ivLength];
				System.arraycopy(data, ivLength, temp, 0, data.length - ivLength);
			} else {
				temp = data;
			}
			decrypt0(temp, stream);
		}
	}

	@Override
	public void decrypt(byte[] data, int length, ByteArrayOutputStream stream) {
		byte[] d = new byte[length];
		System.arraycopy(data, 0, d, 0, length);
		decrypt(d, stream);
	}

	private byte[] randomBytes(int size) {
		byte[] bytes = new byte[size];
		new SecureRandom().nextBytes(bytes);
		return bytes;
	}

	protected abstract StreamBlockCipher getCipher(boolean isEncrypted) throws InvalidAlgorithmParameterException;

	protected abstract void encrypt0(byte[] data, ByteArrayOutputStream stream);

	protected abstract void decrypt0(byte[] data, ByteArrayOutputStream stream);

	protected abstract int getIVLength();

	protected abstract int getKeyLength();
}
